package com.qiangesoft.ratelimit.core.guava;

import com.google.common.util.concurrent.RateLimiter;
import com.qiangesoft.ratelimit.core.LimitKeyBuilder;
import com.qiangesoft.ratelimit.core.LimitType;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Guava限流切面
 *
 * @author qiangesoft
 * @date 2024-03-21
 */
@Slf4j
@Aspect
@Component
public class GuavaLimitAspect {

    private static Map<String, RateLimiter> RATE_HOLDER = new ConcurrentHashMap<>();

    @Before("@annotation(guavaRateLimit)")
    public void doBefore(JoinPoint joinPoint, GuavaRateLimit guavaRateLimit) {
        // 限流逻辑
        LimitType limitType = guavaRateLimit.limitType();
        String key = guavaRateLimit.key();
        int count = guavaRateLimit.count();

        String limitKey = LimitKeyBuilder.build(limitType, key, joinPoint);
        RateLimiter rateLimiter = getRateLimiter(limitKey, count);
        if (!rateLimiter.tryAcquire()) {
            throw new RuntimeException("访问太频繁，请稍后再试！");
        }
    }

    /**
     * 获取RateLimiter
     *
     * @param limitKey
     * @param permitsPerSecond
     * @return
     */
    public synchronized static RateLimiter getRateLimiter(String limitKey, double permitsPerSecond) {
        RateLimiter rateLimiter = RATE_HOLDER.get(limitKey);
        if (rateLimiter != null) {
            return rateLimiter;
        }

        rateLimiter = RateLimiter.create(permitsPerSecond);
        RATE_HOLDER.put(limitKey, rateLimiter);
        return rateLimiter;
    }
}